'''
Created on 19/11/2012

@author: greis
'''
from util import util
import random

class QLearning:
    """
    
    """
    qValue = util.Counter()
    epsilon = 0.5 #probability to take a random action
    discount = 0.5 #the discount rate
    alpha = 0.5 #the learning rate
    
    def getQValue(self, state, action):
        """
          Returns Q(state,action)
          Should return 0.0 if we never seen
          a state or (state,action) tuple
        """
        "*** YOUR CODE HERE ***"
        return self.qValue[(state,action)]
    
    def getValue(self, state):
        """
          Returns max_action Q(state,action)
          where the max is over legal actions.  Note that if
          there are no legal actions, which is the case at the
          terminal state, you should return a value of 0.0.
        """
        "*** YOUR CODE HERE ***"
        legalActions = state.getLegalActions()
        value = []
        for action in legalActions:
            value.append(self.getQValue(state, action))
        if len(value) == 0: 
            return 0.0
        return max(value)

    def getPolicy(self, state):
        """
          Compute the best action to take in a state.  Note that if there
          are no legal actions, which is the case at the terminal state,
          you should return None.
        """
        "*** YOUR CODE HERE ***"
        legalActions = state.getLegalActions()
        bestAction = None
        if len(legalActions) == 0:
            return bestAction
        maxValue = float("-infinity")
        for action in legalActions:
            qValue = self.getQValue(state, action)
            if qValue > maxValue:
                maxValue = qValue
                bestAction = action
            elif qValue == maxValue:
                bestAction = random.choice([action, bestAction])
        return bestAction

    def getAction(self, state):
        """
          Compute the action to take in the current state.  With
          probability self.epsilon, we should take a random action and
          take the best policy action otherwise.  Note that if there are
          no legal actions, which is the case at the terminal state, you
          should choose None as the action.
    
          HINT: You might want to use util.flipCoin(prob)
          HINT: To pick randomly from a list, use random.choice(list)
        """
        # Pick Action
        legalActions = state.getLegalActions()
        action = None
        "*** YOUR CODE HERE ***"
        if len(legalActions) == 0:
            return action
        if util.flipCoin(self.epsilon):
            action = random.choice(legalActions)
        else:
            action = self.getPolicy(state)
        return action

    def update(self, state, action, nextState, reward):
        """
          The parent class calls this to observe a
          state = action => nextState and reward transition.
          You should do your Q-Value update here
    
          NOTE: You should never call this function,
          it will be called on your behalf
        """
        "*** YOUR CODE HERE ***"
        sample = reward + self.discount * self.getValue(nextState)
        self.qValue[state, action] = (1-self.alpha) * self.qValue[state, action] + self.alpha * sample

class ApproximateQLearning(QLearning):
    """
    ApproximateQLearning
    
    You should only have to overwrite getQValue
    and update.  All other QLearning functions
    should work as is.
    """
    
    qValue = util.Counter()
    weightValues = util.Counter()
    
    def getQValue(self, state, action):
        """
          Should return Q(state,action) = w * featureVector
          where * is the dotProduct operator
        """
        "*** YOUR CODE HERE ***"
        self.qValue[(state,action)] = state.getFeatures(action) * self.weightValues
        return self.qValue[(state,action)]

    def update(self, state, action, nextState, reward):
        """
           Should update your weights based on transition
        """
        "*** YOUR CODE HERE ***"
        correction = (reward + self.discount * self.getValue(nextState)) - self.qValue[(state, action)]
        featureValues = state.getFeatures(action)
        for feature in featureValues:
            self.weightValues[feature] = self.weightValues[feature] + self.alpha * correction * featureValues[feature]

